#!/usr/bin/python
# -*- coding: UTF-8 -*-
# code review by phil 21/9/2023 01:52:00
from datetime import datetime, timedelta
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from HV import DailyVol
from dolphindb_data_load import DdbData
from matplotlib.backends.backend_pdf import PdfPages
from scipy.stats import zscore
import matplotlib.gridspec as gridspec
import matplotlib.dates as mdates


class IvAnalysis:
    """
    包括波动率期限结构、两个交易日的波动率对比、偏度数值的波动率计算等、波动率曲面
    """

    def __init__(self):
        self.date = datetime.today().strftime("%Y-%m-%d")
        self.top_months = None  # 最近一个交易日期权总持仓量最大的两个月份

    # 偏度分析
    @staticmethod
    def volatility_term_structure(sym: str = 'TA', months: tuple = ('11', '01')):
        ddb = DdbData()

        # 创建一个n行2列的大图，其中n是syms的长度
        fig, axes = plt.subplots(1, 2, figsize=(15, 6), sharey=True, dpi=500)
        data = ddb.iv_term_for_choose_month(DdbData.s, sym=sym, month=months)

        # 转换日期列为 datetime 类型
        data['date'] = pd.to_datetime(data['date'])

        # 对数据进行排序，确保按日期顺序排列
        data = data.sort_values(by='date')

        # 子图1: 最近5个交易日的 term 的走势
        latest_5_days = data['date'].unique()[-5:]
        data_latest_5_days = data[data['date'].isin(latest_5_days)]
        sns.lineplot(x='date', y='term', data=data_latest_5_days, ax=axes[0])
        axes[0].set_title(f'{months} Term Trend for the Last 5 Trading Days ({sym})')
        axes[0].set_xlabel('Date')
        axes[0].set_ylabel('Term')

        # 子图2: term 的密度分布（垂直于 y 轴）并添加阴影
        sns.kdeplot(y=data['term'], ax=axes[1], fill=True)
        axes[1].set_title(f'{months} Density Distribution of Term with Shade ({sym})')
        axes[1].set_xlabel('Density')
        axes[1].set_ylabel('Term')

        # 确保子图不重叠
        plt.subplots_adjust(hspace=0.2, top=0.85, bottom=0.15)
        plt.suptitle('volatility_term_structure', fontsize=20)  # 添加整体标题
        # plt.savefig('term.jpg')
        # plt.show()

    @staticmethod
    def plot_volatility_cone(end_date: str = '20230922', sym: str = 'TA'):
        """
        将历史波动率锥和指定品种各月份的隐含波动率结合起来，观察隐含波动率的相对位置
        :param end_date:获取主连合约历史数据的结束时间。开始时间为默认的2023年1月1
        :param sym :品种 TA用来从ddb_data获取隐含波动率期限结构数据
        :return:
        """
        # TA0用来获取连续合约历史数据，用来计算历史波动率锥；
        sym0 = sym + '0'
        # 合约列表用来获取历史波动率期限结构
        sym_list = [sym + '2401', sym + '2402', sym + '2405']

        # 拆分列
        def split_underlying(row):
            # 处理合约数据月份，拆分月份和合约名
            alphabetic_part = ''.join(filter(str.isalpha, row['underlying']))
            numeric_part = ''.join(filter(str.isdigit, row['underlying']))
            # Add the omitted digit '202' to the numeric part
            full_year_month = '202' + numeric_part
            return pd.Series([alphabetic_part, full_year_month], index=['symbol', 'month'])

        # 绘图
        plt.figure(figsize=(15, 8), dpi=500)
        hv = DailyVol()
        iv_term = DdbData()
        # 获取历史波动率数据并处理index
        hv_cone_data = hv.calculate_cdf_and_quantiles(hv.vol_complete(hv.data_download(sym0=sym0, e_date=end_date)))
        hv_cone_data = hv_cone_data.reset_index(inplace=False)
        hv_cone_data['index'] = hv_cone_data['index'].astype(int)
        # 获取历史波动率期限结构
        hv_term_data = hv.hv_term(sym_list)
        # 获取隐含波动率数据及到期日，并合并数据，计算期权剩余到期时间
        iv_term_data = iv_term.iv_term_structure_from_dolphindb(DdbData.s, sym)
        latest_date = iv_term_data['date'].max()
        iv_term_data_latest = iv_term_data[iv_term_data['date'] == latest_date]
        iv_term_endday = iv_term.option_end_date(DdbData.s, sym)
        iv_term_endday[['symbol', 'month']] = iv_term_endday.apply(split_underlying, axis=1)
        iv_term_endday_merge = pd.merge(iv_term_data_latest, iv_term_endday, on='month', how='inner')
        iv_term_endday_merge['days_to_expiry'] = (
                iv_term_endday_merge['endday'] - iv_term_endday_merge['date_y']).dt.days
        # 将波动率锥数据、历史波动率期限结构、隐含波动率期限结构数据合并，用于绘图
        iv_hv_merged = pd.merge(iv_term_endday_merge, hv_term_data, on='month')

        # Plot Historical Volatility Cone
        plt.plot(hv_cone_data['index'], hv_cone_data[0.01], label='CDF 0.01', linestyle='-', linewidth=2,
                 color='grey')
        plt.plot(hv_cone_data['index'], hv_cone_data[1], label='CDF 1', linestyle='-', linewidth=2, color='orange')

        # Fill area between 0.25 and 0.01, 1 and 0.75 with a lighter shade of grey
        plt.fill_between(hv_cone_data['index'], hv_cone_data[0.01], hv_cone_data[0.25], color='grey', alpha=0.2,
                         label='Area between CDF 0.01 and 0.25')
        plt.fill_between(hv_cone_data['index'], hv_cone_data[0.75], hv_cone_data[1], color='orange', alpha=0.2,
                         label='Area between CDF 0.75 and 1')

        # Plot Implied Volatility Term Structure
        plt.plot(iv_hv_merged['days_to_expiry'], iv_hv_merged['iv'], color='grey',
                 linestyle='--', marker='o', zorder=5, label='IV now')

        plt.plot(iv_hv_merged['days_to_expiry'], iv_hv_merged['term_hv'], color='grey',
                 linestyle='-', marker='o', zorder=5, label='HV now')

        for a, txt in enumerate(iv_hv_merged['month']):
            plt.annotate(txt, (iv_hv_merged['days_to_expiry'].iloc[a], iv_hv_merged['iv'].iloc[a]),
                         textcoords="offset points", xytext=(0, 10), ha='center', color='grey', alpha=0.3)

        # Labels and title
        plt.xlabel('Days to Expiry')
        plt.ylabel('Volatility')
        plt.title(f'Volatility Cone with IV&HV Term Structure for {sym}')
        plt.legend()
        plt.grid(True)

        plt.subplots_adjust(hspace=0.2, top=0.90, bottom=0.1)
        plt.suptitle('Volatility', fontsize=20)  # 添加整体标题
        # plt.savefig('vol_cone.jpg')

    @staticmethod
    def plot_hv_iv_diff_volatility(sym: str = 'TA', combinations=None):
        """
        获取并计算20日历史波动率数据和隐含波动率数据，并分析历史波动率和隐含波动率的价差密度分布、最近14个交易日的价差走势、历史波动率和隐含波动率的
        走势图
        :param sym:品种 TA\MA\RM\CF
        :param combinations: 参数组合(full_sym:用来获取新浪财经指定期货合约数据，sym,month:两个参数一起用来从ddb查询隐含波动率，symbol：用来在ddb分钟级历史行情
        :return:
        """
        if combinations is None:
            combinations = [(sym + '2401', sym, '202401', sym + '401'), (sym + '2405', sym, '202405', sym + '405')]

        n = len(combinations)  # Number of subplots needed
        fig, axes = plt.subplots(n, 3, figsize=(15, 10 * n), dpi=500)

        for i, (full_sym, sym, month, symbol) in enumerate(combinations):
            hv = DailyVol()
            ddb = DdbData()

            data_hv = hv.hv(hv.data_download(main=False, symbol=full_sym), method='hv_abs')
            data_iv = ddb.fetch_iv_from_dolphindb(DdbData.s, sym, month)
            # data_rv = hv.rv(ddb.minute_data_load(DdbData.s, symbol))
            # data_rv['date'] = data_rv['date'].astype(str)

            tb = pd.merge(data_hv, data_iv, on=['date'], how='inner')
            tb.columns = ['date', 'data_hv', 'data_iv']
            tb['hv_iv_diff'] = tb['data_hv'] - tb['data_iv']
            tb['date'] = tb['date'].astype(str)
            #
            # Make the first two columns share the y-axis
            axes[i, 0].get_shared_y_axes().join(axes[i, 0], axes[i, 1])

            # Plotting the last 5 days
            last_days = tb.tail(14)
            sns.lineplot(x='date', y='hv_iv_diff', data=last_days, marker='o', ax=axes[i, 0])
            axes[i, 0].set_title(f'hv_iv_diff Trend for the Last 14 Trading Days {full_sym}')
            axes[i, 0].set_xticklabels(tb['date'].tail(14), rotation=30)
            axes[i, 0].tick_params(axis='x', labelsize=8)  # 可以调整标签大小
            axes[i, 0].set_ylabel('hv_iv_diff')

            # Plotting the density
            sns.kdeplot(tb['hv_iv_diff'], label='Density of hv_iv_diff', fill=True, ax=axes[i, 1], vertical=True)
            axes[i, 1].set_title(f'Density Distribution of hv_iv_diff {full_sym}')
            axes[i, 1].set_xlabel('Density')

            # Third subplot: Lineplot of 'hv_abs' and 'iv'
            sns.lineplot(x='date', y='data_hv', data=tb, label='hv', ax=axes[i, 2])
            sns.lineplot(x='date', y='data_iv', data=tb, label='iv', ax=axes[i, 2])
            axes[i, 2].set_title(f'Historical Trends of HV and IV for {full_sym}')
            all_dates = tb['date']
            # 设置间隔 - 例如，每5天显示一个标签
            interval = 20
            # 计算间隔的日期索引，确保包括最后一个日期
            date_indices = list(range(0, len(all_dates), interval))
            # 确保最后一个日期（最新的日期）始终显示在图表的最右侧
            if date_indices[-1] != len(all_dates) - 1:
                date_indices.append(len(all_dates) - 1)
            # 获取要显示的日期标签
            dates_to_show = [all_dates[i] for i in date_indices]
            # 应用日期索引和标签到图表
            axes[i, 2].set_xticks(date_indices)
            axes[i, 2].set_xticklabels(dates_to_show, rotation=45)  # 旋转标签以更好地适应，角度可以调整
            axes[i, 2].tick_params(axis='x', labelsize=8)  # 可以调整标签大小
            axes[i, 2].set_xlabel('Date')
            axes[i, 2].set_ylabel('Volatility')
            axes[i, 2].legend()
            axes[i, 2].set_ylim(0.05, 0.5)

        plt.suptitle('HV&IV Volatility Analysis for ' + sym, fontsize=20)  # 添加整体标题
        plt.subplots_adjust(hspace=0.2, top=0.9, bottom=0.1)

    def plot_rv_iv(self, sym: str = 'TA'):
        """
        从ddb获取分钟级行情数据和收盘隐含波动率数据，用于计算实现波动率日度年化波动率，并于收盘结算iv做比较。
        展示持仓量最大的两个月份
        :param sym: 选定的品种 TA 获取iv
        :return:
        """
        # 实例化
        realized_vol = DailyVol()
        implied_vol = DdbData()

        fig, axes = plt.subplots(1, 2, figsize=(15, 8))  # 1 row, 2 columns

        for i, month_suffix in enumerate(self.top_months):
            year = f"2024{month_suffix}"  # Adjust year
            symbol = f"{sym}4{month_suffix}"  # Adjust symbol

            # Calculate realized volatility
            rv_data = realized_vol.rv(implied_vol.minute_data_load(DdbData.s, symbol))

            # Fetch implied volatility
            iv_data = implied_vol.fetch_iv_from_dolphindb(DdbData.s, sym, year)

            # Data preprocessing
            rv_data['date'] = pd.to_datetime(rv_data['date'])
            plot_data = pd.merge(rv_data, iv_data, on='date', how='inner').dropna()

            # Plotting in subplot
            ax = axes[i]
            ax.plot(plot_data['date'], plot_data['rv'], label='Realized Volatility')
            ax.plot(plot_data['date'], plot_data['iv'], label='Implied Volatility')
            ax.set_xlabel('Date')
            ax.set_ylabel('Volatility')
            ax.set_title(f'{symbol}')
            ax.legend()
            ax.tick_params(axis='x', rotation=45)
        plt.title(f'Realized vs Implied Volatility for {sym}', fontsize=20)
        plt.subplots_adjust(hspace=0.2, top=0.9, bottom=0.15)


    @staticmethod
    def plot_iv_skew(sym: str = 'TA', months: tuple = ('01', '05')):
        """
        绘制指定品种，指定月份的偏度绘制
        :param sym: 指定品种
        :param months: 指定月份
        :return: 偏度分布图
        """
        ddb = DdbData()
        # Assuming the length of months is 2 for a 2x2 grid layout
        gs = gridspec.GridSpec(2, 2, width_ratios=[2, 0.5], hspace=0.5, wspace=0.2)
        fig = plt.figure(figsize=(15, 8), dpi=500)

        for i, month in enumerate(months):
            data = ddb.iv_skew_all(DdbData.s, sym=sym, month_new=month)
            if data.empty:
                continue

            data['date'] = pd.to_datetime(data['date'])
            date_format = mdates.DateFormatter('%Y-%m-%d')
            data['z_score'] = data.groupby('month_new')['skew'].transform(lambda x: zscore(x, nan_policy='omit'))
            data_for_densityplot = data[np.abs(data['z_score']) < 6]
            data_for_lineplot = data[data['date'] > data['date'].max() - pd.Timedelta(days=5)]

            # Line plot
            ax1 = fig.add_subplot(gs[i, 0])
            sns.lineplot(data=data_for_lineplot, x='date', y='skew', ax=ax1)
            ax1.xaxis.set_major_formatter(date_format)
            ax1.xaxis.set_major_locator(mdates.DayLocator(interval=1))
            ax1.set_title(f'Skewness for {sym} in {month} (Last 5 Days)')
            ax1.grid(True, linestyle='--', alpha=0.6)

            # Density plot
            ax2 = fig.add_subplot(gs[i, 1], sharey=ax1)
            sns.kdeplot(y=data_for_densityplot['skew'], ax=ax2, fill=True)
            ax2.set_xlabel('Density')
            ax2.set_ylabel('')

        plt.subplots_adjust(hspace=0.2, top=0.85, bottom=0.15)
        plt.suptitle('IV Skew Analysis for ' + sym, fontsize=20)

    @staticmethod
    def list_top_change_contract_for_vol_int():
        """
        展示所有期权合约中成交量和持仓量变化最大的前4个合约
        :return:
        """
        ddb = DdbData()
        top_df_vol = ddb.top_vol_int_option(DdbData.s, 'diff_vol')
        top_df_int = ddb.top_vol_int_option(DdbData.s, 'diff_int')

        # 创建图表
        figs = []
        # 绘制 top_df_vol
        fig_vol, ax_vol = plt.subplots(figsize=(15, 8))
        ax_vol.axis('off')
        ax_vol.axis('tight')
        cell_text = top_df_vol.values.tolist()  # 转换 DataFrame 为二维列表
        col_labels = top_df_vol.columns.tolist()  # 获取列标签
        tbl = ax_vol.table(cellText=cell_text, colLabels=col_labels, loc='center')  # 创建表格
        tbl.auto_set_font_size(False)
        tbl.set_fontsize(10)  # 设置字体大小
        tbl.scale(1.2, 1.2)  # 缩放表格
        plt.suptitle('table for top 4 changed contract for volume', fontsize=20)
        figs.append(fig_vol)
        # 绘制 top_df_int
        fig_int, ax_int = plt.subplots(figsize=(15, 8))
        ax_int.axis('off')
        ax_int.axis('tight')
        cell_text = top_df_int.values.tolist()  # 转换 DataFrame 为二维列表
        col_labels = top_df_int.columns.tolist()  # 获取列标签
        tbl = ax_int.table(cellText=cell_text, colLabels=col_labels, loc='center')  # 创建表格
        tbl.auto_set_font_size(False)
        tbl.set_fontsize(10)  # 设置字体大小
        tbl.scale(1.2, 1.2)  # 缩放表格
        plt.suptitle('table for top 4 changed contract for int', fontsize=20)
        figs.append(fig_int)

        plt.subplots_adjust(hspace=0.2, top=0.85, bottom=0.15)
        return figs

    def polt_option_volume_open_interest(self, sym: str):
        """
        从ddb获取期权的成交量和持仓量的日度的历史数据，并绘制面积图，绘制1行两列图表，将最新一个交易日的持仓量最大的两个月份的
        分别的看涨和看跌期权的持仓量绘制在左边，将成交量绘制在右边，以时间为x轴。
        :param sym: 选定的品种
        :return:
        """
        ddb = DdbData()
        plt.style.use('seaborn')
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 8))

        for month in self.top_months:
            for option_type in ['C', 'P']:
                data = ddb.option_data_volume_openinterest(DdbData.s, sym, month, option_type)
                if data.empty:
                    continue

                data['date'] = pd.to_datetime(data['date'])
                data.sort_values(by='date', inplace=True)

                # 筛选最近一个月的数据
                latest_date = data['date'].max()
                one_month_ago = latest_date - timedelta(days=30)
                recent_data = data[data['date'] >= one_month_ago]

                # 左图：持仓量
                ax1.plot(recent_data['date'], recent_data['sum_openinterest'],
                         label=f'{month} {option_type}')
                # 右图：成交量
                ax2.plot(recent_data['date'], recent_data['sum_volume'],
                         label=f'{month} {option_type}')

        # 设置 X 轴刻度以确保显示最新日期
        for ax in [ax1, ax2]:
            ax.set_xlim(one_month_ago, latest_date)
            ax.xaxis.set_major_locator(mdates.AutoDateLocator(minticks=10, maxticks=10))
            ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))

        # 设置图表标题和轴标签
        ax1.set_title(f'Open_Interest for {sym}', fontsize=15)
        ax2.set_title(f'Volume for {sym}', fontsize=15)
        ax1.set_xlabel('Date')
        ax2.set_xlabel('Date')
        ax1.set_ylabel('Open Interest')
        ax2.set_ylabel('Volume')
        ax1.grid(True)
        ax2.grid(True)

        # 添加图例
        ax1.legend()
        ax2.legend()

        plt.subplots_adjust(hspace=0.2, top=0.85, bottom=0.15)
        plt.suptitle(f'open_interset&int for {sym}', fontsize=20)  # 添加整体标题

    def plot_put_call_ratio(self, sym: str):
        """
        从ddb获取期权数据，并绘制最近一天持仓量和成交量最大的两个月份的最近一个月的持仓量和成交量pcr
        :param sym: 选定的品种
        :return: pcr图
        """
        # Assuming ddb.get_data() fetches the required data
        ddb = DdbData()
        data = ddb.put_call_ratio(DdbData.s, sym)

        # Calculate total open interest and determine top two months
        data['total_open_interest'] = data['sum_openinterest'] + data['data_put_sum_openinterest']
        latest_date = data['date'].max()
        # 找到最近一个月的时间范围
        one_month_ago = latest_date - pd.Timedelta(days=60)
        # 筛选出最近一个月的数据
        recent_data = data[data['date'] > one_month_ago]

        # 筛选出最近日期的数据
        latest_data = data[data['date'] == latest_date]
        self.top_months = latest_data.groupby('month')['total_open_interest'].sum().nlargest(2).index.tolist()

        fig, axes = plt.subplots(1, 2, figsize=(15, 8))
        for i, month in enumerate(self.top_months):
            month_data = recent_data[recent_data['month'] == month]
            month_data['date'] = pd.to_datetime(month_data['date'])
            sns.lineplot(data=month_data, x='date', y='vol_pcr', ax=axes[i], label='Volume', ci=None)
            sns.lineplot(data=month_data, x='date', y='interest_pcr', ax=axes[i], label='Open_Interest', ci=None)
            # 设置x轴日期格式和标签
            axes[i].xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))

            # 计算合适的间隔
            total_days = (month_data['date'].max() - month_data['date'].min()).days
            interval = max(1, total_days // 6)  # 确保至少显示6个标签

            # 设置标签间隔，确保最新日期显示
            axes[i].xaxis.set_major_locator(mdates.DayLocator(interval=interval))
            axes[i].tick_params(axis='x', rotation=45)  # 旋转45度

            axes[i].set_title(f'{sym} - {month}')
            axes[i].legend()

        plt.suptitle(f'Volume and Interest PCR for {sym}', fontsize=20)
        plt.subplots_adjust(hspace=0.2, top=0.85, bottom=0.15)

    def generate_pdf_with_matplotlib(self, end_date, output_pdf_name: str = None):
        """
        使用Matplotlib的PdfPages生成PDF。
        :param output_pdf_name: pdf名称
        :param end_date:获取主连合约历史数据的结束时间。开始时间为默认的2023年1月1
        """

        if output_pdf_name is None:
            output_pdf_name = self.date + "IV_HV_Analysis.pdf"

        # 使用PdfPages保存多个图像为一个PDF文件
        with PdfPages(output_pdf_name) as pdf:
            for i in ['TA', 'MA', 'RM', 'CF']:
                # todo 将持仓量最大的月份判断参数全局化，让其他几个函数不在需要指定月份，而是按照持仓量最大的两个月份直接来画图。
                self.plot_put_call_ratio(sym=i)
                pdf.savefig()
                plt.close()

                self.plot_hv_iv_diff_volatility(sym=i)
                pdf.savefig()
                plt.close()

                self.plot_volatility_cone(end_date=end_date, sym=i)
                pdf.savefig()
                plt.close()

                self.volatility_term_structure(sym=i, months=('01', '05'))
                pdf.savefig()  # 保存当前图像为PDF的一个页面
                plt.close()  # 关闭当前图像

                self.plot_iv_skew(sym=i, months=('01', '05'))
                pdf.savefig()
                plt.close()

                self.polt_option_volume_open_interest(sym=i)
                pdf.savefig()
                plt.close()

                self.plot_rv_iv(sym=i)
                pdf.savefig()
                plt.close()

            # 获取包含表格的图形对象
            figs = self.list_top_change_contract_for_vol_int()
            for fig in figs:
                pdf.savefig(fig, bbox_inches='tight')
                plt.close(fig)


if __name__ == '__main__':
    iv = IvAnalysis()
    iv.generate_pdf_with_matplotlib(end_date='20231128')

